import torch
import prob

torch.set_default_dtype(torch.float64)
torch.set_default_tensor_type(torch.DoubleTensor)

############################################
# constraints for barrier certificate B:
# (1) init ==> B <= 0
# (2) unsafe ==> B > 0
# (3) domain ==>   B(x') - (1-alpha) * B(x) <= 0
############################################

def calc_loss(model, input_init, input_unsafe, input_domain):
    output_domain = model(input_domain)  # B(x)

    # loss of domain
    if len(input_domain) == 0:
        loss_domain = torch.tensor(0) * output_domain
        # avg_loss_domain = torch.tensor(0) * output_domain[0, 0]
    else:
        vector_domain = prob.vector_field(input_domain)  # x'
        output_vector_domain = model(vector_domain)  # B(x')
        condition_domain = output_vector_domain - (1 - prob.alpha) * output_domain  # B(x') - (1 - alpha) * B(x)
        loss_domain = torch.relu(condition_domain + prob.TOL_DOMAIN)
        # avg_loss_domain = torch.sum(loss_domain) / loss_domain.size()[0]

    # loss of init
    if len(input_init) == 0:
        loss_init = torch.tensor(0) * output_domain
        # avg_loss_init = torch.tensor(0) * output_domain[0, 0]
    else:
        output_init = model(input_init)
        loss_init = torch.relu(output_init + prob.TOL_INIT)
        # avg_loss_init = torch.sum(loss_init) / loss_init.size()[0]

    # loss of unsafe
    if len(input_unsafe) == 0:
        loss_unsafe = torch.tensor(0) * output_domain
        # avg_loss_unsafe = torch.tensor(0) * output_domain[0, 0]
    else:
        output_unsafe = model(input_unsafe)
        loss_unsafe = torch.relu((- output_unsafe) + prob.TOL_UNSAFE)
        # avg_loss_unsafe = torch.sum(loss_unsafe) / loss_unsafe.size()[0]

    loss_init_sum = prob.DECAY_INIT * torch.sum(loss_init)
    loss_unsafe_sum = prob.DECAY_UNSAFE * torch.sum(loss_unsafe)
    loss_domain_sum = prob.DECAY_DOMAIN * torch.sum(loss_domain)

    total_loss = loss_init_sum + loss_unsafe_sum + loss_domain_sum

    # return total_loss
    return total_loss, loss_init_sum, loss_unsafe_sum, loss_domain_sum
